'''
$Revision: 124 $
$LastChangedDate: 2009-05-06 14:29:18 +0000 (Wed, 06 May 2009) $
'''

import unittest
import os, sys
from sympy import Symbol, sympify, diff


from MomentClosure.GetMomentEquation import getMomentEquation
from MomentClosure.GetReactionInformation import GetStoichiometryDict
from MomentClosure.VariableDict import Variables

from MomentClosure.GlobalVariables import globalVariables
from MomentClosure.GetTimeDerivatives import getTimeDerivative


class TestMomentEquations(unittest.TestCase):
    '''Tests the moment closure equations with know cases.
    
        If they work, everything else should as well.
    '''
    def setUp(self):
        pass

    def testImmigrationDeath(self):
        Variables().initialise()
        self.__testImmigrationDeath('xml_files/test/imdeath.xml')
        #With local parameter
        self.__testImmigrationDeath('xml_files/test/ch05-imdeath.xml')
        
        #Reactions of the same type
        self.__testImmigrationDeath('xml_files/test/imdeath2.xml')
    
    def __testImmigrationDeath(self, file_name):
        Variables().initialise()
        gdd = GetStoichiometryDict(sbml_file = file_name)
        reaction_list = gdd.getReactionList()
        
        l = Symbol('lambda1')
        nu = Symbol('mu')
        t= Symbol('t')
        t_der = diff(sympify('X(t)'), t)
        t_2der = diff(sympify('E_X_X(t)'), t)
        t_3der = diff(sympify('mu_3(t)'), t)
        
        
        mu1 = Variables().getMoment((1,))
        mu2 = Variables().getMoment((2,))
        mu3 = Variables().getMoment((3,)) 
       
        
        x=getTimeDerivative((0,), gdd)
        mom_eq = getMomentEquation((0,), reaction_list, gdd.hasDenominator) 
        self.assertEqual(mom_eq, 0)  
               
        #Mean
        mom_eq = getMomentEquation((1,), reaction_list, gdd.hasDenominator) 
        x = getTimeDerivative((1,), gdd)
        exact = l - nu*mu1
        difference = mom_eq/x.coeff(t_der) - exact
        self.assertEqual(difference.expand(), 0)  

        #Mu[2]
        mom_eq = getMomentEquation((2,), reaction_list, gdd.hasDenominator)
        x=getTimeDerivative((2,), gdd)
        exact = l*(1+2*mu1)+nu*(mu1-2*mu2)
        difference = mom_eq/x.coeff(t_2der) - exact
        self.assertEqual(difference.expand(), 0)  
        
        #Mu[3]        
        mom_eq = getMomentEquation((3,), reaction_list, gdd.hasDenominator)
        x=getTimeDerivative((3,), gdd)
        exact = l*(1+3*mu1+3*mu2) - nu*(mu1-3*mu2+3*mu3)
        difference = mom_eq/x.coeff(t_3der) - exact
        self.assertEqual(difference.expand(), 0) 


    def testDimerModel(self):
        Variables().initialise()
        
        gdd = GetStoichiometryDict(sbml_file = 'xml_files/smfsb/single-dimer.xml')
        reaction_list = gdd.getReactionList()
        
        k1 = Symbol('k1')
        k2 = Symbol('k2')
        P0 = Symbol('P0')
        t= Symbol('t')
        t_der = diff(sympify('P(t)'), t)
        t_2der = diff(sympify('E_P_P(t)'), t)
        t_3der = diff(sympify('mu_3(t)'), t)

        mu1 = Variables().getMoment((1,))
        mu2 = Variables().getMoment((2,))
        mu3 = Variables().getMoment((3,)) 
        mu4 = Variables().getMoment((4,))
        
        mom_eq = getMomentEquation((0,), reaction_list, gdd.hasDenominator) 
        self.assertEqual(mom_eq, 0)
                  
        mom_eq = getMomentEquation((1,), reaction_list, gdd.hasDenominator) 
        x = getTimeDerivative((1,), gdd)
        exact = k1*(mu1-mu2)+k2*(P0-mu1)
        difference = mom_eq/x.coeff(t_der) - exact
        self.assertEqual(difference.expand(), 0)  
        
        mom_eq = getMomentEquation((2,), reaction_list, gdd.hasDenominator)
        x = getTimeDerivative((2,), gdd)
        exact = -2*k1*(mu1-2*mu2+mu3)+2*k2*(P0*(mu1+1)-mu1-mu2)
        difference = mom_eq/x.coeff(t_2der) - exact
        self.assertEqual(difference.expand(), 0)  

        mom_eq = getMomentEquation((3,), reaction_list, gdd.hasDenominator)
        x = getTimeDerivative((3,), gdd)
        exact = k1*(4*mu1-10*mu2+9*mu3-3*mu4)
        exact += k2*(4*(P0-mu1)+6*(P0*mu1-mu2)+3*(P0*mu2-mu3))
        difference = mom_eq/x.coeff(t_3der) - exact
        self.assertEqual(difference.expand(), 0)  


    def testSIS(self):
        Variables().initialise()
        
        gdd = GetStoichiometryDict(sbml_file = 'xml_files/sis.xml')
        reaction_list = gdd.getReactionList()
        
        alpha = Symbol('alpha')
        mu = Symbol('mu')
        beta = Symbol('beta')
        delta = Symbol('delta')
        
        mu10 = Variables().getMoment((1,0))
        mu01 = Variables().getMoment((0,1))
        
        mu20 = Variables().getMoment((2,0)) 
        mu11 = Variables().getMoment((1,1)) 
        mu02 = Variables().getMoment((0,2))
        
        mu30 = Variables().getMoment((3,0)) 
        mu21 = Variables().getMoment((2,1)) 
        mu12 = Variables().getMoment((1,2)) 
        mu03 = Variables().getMoment((0,3))        
        
        mom_eq = getMomentEquation((0,0), reaction_list, gdd.hasDenominator) 
        self.assertEqual(mom_eq, 0)  
               
        #Mean 10
        mom_eq = getMomentEquation((1,0), reaction_list, gdd.hasDenominator)
        exact = alpha - beta*mu11
        difference = mom_eq - exact
        self.assertEqual(difference, 0)  

        #Mean 01
        mom_eq = getMomentEquation((0,1), reaction_list, gdd.hasDenominator)
        exact = delta + beta*mu11 - mu*mu01
        difference = mom_eq - exact
        self.assertEqual(difference, 0)  

        #mu20
        mom_eq = getMomentEquation((2,0), reaction_list, gdd.hasDenominator)
        exact = alpha*(1+2*mu10) + beta*mu11-2*beta*mu21
        difference = mom_eq - exact
        self.assertEqual(difference.expand(), 0)  

        #mu11
        mom_eq = getMomentEquation((1,1), reaction_list, gdd.hasDenominator)
        exact = beta*(mu21-mu11-mu12)-mu*mu11+delta*mu10+alpha*mu01
        difference = mom_eq - exact
        self.assertEqual(difference.expand(), 0)  
        
        #mu02
        mom_eq = getMomentEquation((0,2), reaction_list, gdd.hasDenominator)
        exact = beta*mu11+2*beta*mu12+mu*mu01-2*mu*mu02+delta*(1+2*mu01)
        difference = mom_eq - exact
        self.assertEqual(difference.expand(), 0)  
        
    def testDummyModel(self):
        
        Variables().initialise()
        globalVariables(closure_order = 4)
        gdd = GetStoichiometryDict(sbml_file = 'xml_files/testReaction.xml')
        reaction_list = gdd.getReactionList()
       
        mom_eq = getMomentEquation((0,0,0,0,0), reaction_list, gdd.hasDenominator) 
        self.assertEqual(mom_eq, 0)  
               
        l = Symbol('lambda1')
        m = Symbol('mu')

        mu10000 = Variables().getMoment((1,0,0,0,0))
        mu01000 = Variables().getMoment((0,1,0,0,0))        
        mu00100 = Variables().getMoment((0,0,1,0,0))
        mu00010 = Variables().getMoment((0,0,0,1,0))
        mu00001 = Variables().getMoment((0,0,0,0,1))
        

        mu11000 = Variables().getMoment((1,1,0,0,0))
#        mu01000 = Variables().getMoment((0,1,0,0,0))        
#        mu00100 = Variables().getMoment((0,0,1,0,0))
#        mu00010 = Variables().getMoment((0,0,0,1,0))
#        mu00001 = Variables().getMoment((0,0,0,0,1))

                                

#        exact = alpha - beta*mu11
#        difference = mom_eq - exact
#        self.assertEqual(difference, 0)  

        
        
        
                        
if __name__ == '__main__':
    unittest.main()




























